{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bf8e01e7-df86-450d-8f01-6573320a93b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyscf import gto, scf, cc, lib\n",
    "import numpy as np\n",
    "from functools import partial\n",
    "import os\n",
    "\n",
    "np.einsum = partial(np.einsum, optimize=True)\n",
    "np.set_printoptions(6, suppress=True, linewidth=150)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "758eb78f-632c-4f32-82b9-712687fc8d8f",
   "metadata": {},
   "source": [
    "## System Definition and Reference CC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "47412686-c36d-4e0d-bca6-2be749ef60f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 10.1021/jp104865w, (H2O)10, PP5 structure\n",
    "mol = gto.Mole(atom=\"\"\"\n",
    "    O -2.21165 0.99428 -1.34761\n",
    "    H -1.39146 1.51606 -1.47747\n",
    "    H -1.97320 0.08049 -1.61809\n",
    "    O 0.09403 2.29278 1.59474\n",
    "    H 0.12603 2.53877 0.64902\n",
    "    H -0.74393 1.78978 1.67135\n",
    "    O -1.36387 -1.68942 -1.58413\n",
    "    H -1.87986 -2.36904 -2.04608\n",
    "    H -1.51808 -1.85775 -0.60321\n",
    "    O 1.15753 -1.98493 1.42883\n",
    "    H 1.51336 -1.05256 1.56992\n",
    "    H 1.63126 -2.54067 2.06706\n",
    "    O 2.16234 0.46384 1.59959\n",
    "    H 1.45220 1.14162 1.73767\n",
    "    H 2.44819 0.61600 0.67631\n",
    "    O 0.26320 2.39844 -1.29615\n",
    "    H 1.04651 1.79827 -1.38236\n",
    "    H 0.46651 3.18119 -1.83082\n",
    "    O 1.44377 -1.86519 -1.36370\n",
    "    H 0.48945 -1.86011 -1.60072\n",
    "    H 1.44320 -2.10978 -0.41122\n",
    "    O -1.62831 -1.98091 1.04938\n",
    "    H -1.92768 -1.08892 1.33229\n",
    "    H -0.69028 -2.03600 1.33896\n",
    "    O 2.35473 0.62384 -1.26848\n",
    "    H 3.15897 0.65726 -1.80967\n",
    "    H 2.00663 -0.31760 -1.36507\n",
    "    O -2.29362 0.74293 1.32406\n",
    "    H -2.34790 0.87628 0.33220\n",
    "    H -3.13510 1.07144 1.67759\n",
    "\"\"\", basis=\"cc-pVDZ\").build()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "623e5681-fa5c-40dd-a44b-8974bfba0f23",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "converged SCF energy = -760.417170870121\n"
     ]
    }
   ],
   "source": [
    "mf = scf.RHF(mol).density_fit(auxbasis=\"cc-pVDZ-ri\").run()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc1062e1-0373-4c20-8e9d-49dffe11af24",
   "metadata": {},
   "outputs": [],
   "source": [
    "mol.verbose = 4\n",
    "lib.logger.TIMER_LEVEL = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c49827cb-8a1d-42c6-a969-456d8983f67d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "******** <class 'pyscf.cc.dfccsd.RCCSD'> ********\n",
      "CC2 = 0\n",
      "CCSD nocc = 40, nmo = 230\n",
      "frozen orbitals 10\n",
      "max_cycle = 50\n",
      "direct = 0\n",
      "conv_tol = 1e-07\n",
      "conv_tol_normt = 1e-05\n",
      "diis_space = 6\n",
      "diis_start_cycle = 0\n",
      "diis_start_energy_diff = 1e+09\n",
      "max_memory 4000 MB (current use 450 MB)\n",
      "    CPU time for df vj and vk      0.55 sec, wall time      0.03 sec\n",
      "    CPU time for df vj and vk      0.57 sec, wall time      0.04 sec\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<class 'pyscf.cc.dfccsd.RCCSD'> does not have attributes  converged\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Init t2, MP2 energy = -762.512442320864  E_corr(MP2) -2.09527145074345\n",
      "    CPU time for init mp2      2.41 sec, wall time      0.41 sec\n",
      "Init E_corr(RCCSD) = -2.09527145074912\n",
      "cycle = 1  E_corr(RCCSD) = -2.12370362944668  dE = -0.0284321787  norm(t1,t2) = 0.0841487\n",
      "    CPU time for RCCSD iter    396.32 sec, wall time     26.83 sec\n",
      "cycle = 2  E_corr(RCCSD) = -2.16249030073982  dE = -0.0387866713  norm(t1,t2) = 0.0246926\n",
      "    CPU time for RCCSD iter    439.82 sec, wall time     28.49 sec\n",
      "cycle = 3  E_corr(RCCSD) = -2.16887833109026  dE = -0.00638803035  norm(t1,t2) = 0.0115431\n",
      "    CPU time for RCCSD iter    455.28 sec, wall time     29.67 sec\n",
      "cycle = 4  E_corr(RCCSD) = -2.17377752172915  dE = -0.00489919064  norm(t1,t2) = 0.00379069\n",
      "    CPU time for RCCSD iter    447.13 sec, wall time     29.23 sec\n",
      "cycle = 5  E_corr(RCCSD) = -2.17369713739983  dE = 8.03843293e-05  norm(t1,t2) = 0.000948251\n",
      "    CPU time for RCCSD iter    448.51 sec, wall time     29.29 sec\n",
      "cycle = 6  E_corr(RCCSD) = -2.17354870722074  dE = 0.000148430179  norm(t1,t2) = 0.000314353\n",
      "    CPU time for RCCSD iter    450.83 sec, wall time     29.61 sec\n",
      "cycle = 7  E_corr(RCCSD) = -2.17354099904743  dE = 7.70817331e-06  norm(t1,t2) = 0.000117935\n",
      "    CPU time for RCCSD iter    452.58 sec, wall time     29.94 sec\n",
      "cycle = 8  E_corr(RCCSD) = -2.17355086802812  dE = -9.8689807e-06  norm(t1,t2) = 3.79795e-05\n",
      "    CPU time for RCCSD iter    454.12 sec, wall time     29.90 sec\n",
      "cycle = 9  E_corr(RCCSD) = -2.17354991457266  dE = 9.5345546e-07  norm(t1,t2) = 1.4716e-05\n",
      "    CPU time for RCCSD iter    458.24 sec, wall time     30.23 sec\n",
      "cycle = 10  E_corr(RCCSD) = -2.17354987210921  dE = 4.24634581e-08  norm(t1,t2) = 5.22321e-06\n",
      "    CPU time for RCCSD iter    451.58 sec, wall time     29.85 sec\n",
      "    CPU time for RCCSD   4454.42 sec, wall time    293.03 sec\n",
      "RCCSD converged\n",
      "E(RCCSD) = -762.5907207422301  E_corr = -2.173549872109206\n"
     ]
    }
   ],
   "source": [
    "mf_cc = cc.CCSD(mf, frozen=10).run()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ee93f1f-9950-4f37-8595-92a539fa2842",
   "metadata": {},
   "source": [
    "## Global Variables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4c70ae44-6571-4917-a36d-e755350acedb",
   "metadata": {},
   "outputs": [],
   "source": [
    "nfrz = 10\n",
    "cderi_ao = mf.with_df._cderi\n",
    "mo_occ = mf.mo_occ[nfrz:]\n",
    "mo_energy = mf.mo_energy[nfrz:]\n",
    "mo_coeff = mf.mo_coeff[:, nfrz:]\n",
    "\n",
    "nmo = len(mo_occ)\n",
    "nocc = (mo_occ > 0).sum()\n",
    "nvir = nmo - nocc\n",
    "nao = mol.nao\n",
    "naux = cderi_ao.shape[0]\n",
    "\n",
    "so = slice(0, nocc)\n",
    "sv = slice(nocc, nmo)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "45dd19e3-3201-46f7-990b-cb1012dc60cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "B = np.einsum(\"Puv, up, vq -> pqP\", lib.unpack_tril(cderi_ao), mo_coeff, mo_coeff)\n",
    "D_ov = mo_energy[so, None] - mo_energy[None, sv]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "bd4dd9fb-4547-4e0d-8a36-d2e3ce8e9695",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.makedirs(\"h2o_pp5-cc-pvdz\", exist_ok=True)\n",
    "np.save(\"h2o_pp5-cc-pvdz/mo_coeff.npy\", np.ascontiguousarray(mo_coeff))\n",
    "np.save(\"h2o_pp5-cc-pvdz/mo_occ.npy\", mo_occ)\n",
    "np.save(\"h2o_pp5-cc-pvdz/mo_energy.npy\", mo_energy)\n",
    "np.save(\"h2o_pp5-cc-pvdz/cderi.npy\", mf.with_df._cderi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "788502d5-b4d4-428f-99b6-a6cc82a6c995",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
